/*
 *	Copyright 2023 Jan Pfeifer
 *
 *	Licensed under the Apache License, Version 2.0 (the "License");
 *	you may not use this file except in compliance with the License.
 *	You may obtain a copy of the License at
 *
 *	http://www.apache.org/licenses/LICENSE-2.0
 *
 *	Unless required by applicable law or agreed to in writing, software
 *	distributed under the License is distributed on an "AS IS" BASIS,
 *	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *	See the License for the specific language governing permissions and
 *	limitations under the License.
 */

package graph

import (
	"fmt"
	"reflect"
	"runtime"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/gomlx/gomlx/backends"
	"github.com/gomlx/gomlx/internal/exceptions"
	"github.com/gomlx/gomlx/pkg/core/distributed"
	"github.com/gomlx/gomlx/pkg/core/shapes"
	"github.com/gomlx/gomlx/pkg/core/tensors"
	"github.com/gomlx/gomlx/pkg/support/xslices"
	"github.com/pkg/errors"
	"k8s.io/klog/v2"
)

// Generated by `cmd/constraints_generator`:

// ExecGraphFn is a type parameter for accepted function types for MustNewExec constructor.
//
//nolint:goimports,golines // See https://youtrack.jetbrains.com/issue/GO-19556/Goland-formatter-gofmt
type ExecGraphFn interface {
	func(*Graph) *Node |
		func([]*Node) *Node |
		func(*Node) *Node |
		func(*Node, *Node) *Node |
		func(*Node, *Node, *Node) *Node |
		func(*Node, *Node, *Node, *Node) *Node |
		func(*Node, *Node, *Node, *Node, *Node) *Node |
		func(*Node, *Node, *Node, *Node, *Node, *Node) *Node |
		func(*Graph) (*Node, *Node) |
		func([]*Node) (*Node, *Node) |
		func(*Node) (*Node, *Node) |
		func(*Node, *Node) (*Node, *Node) |
		func(*Node, *Node, *Node) (*Node, *Node) |
		func(*Node, *Node, *Node, *Node) (*Node, *Node) |
		func(*Node, *Node, *Node, *Node, *Node) (*Node, *Node) |
		func(*Node, *Node, *Node, *Node, *Node, *Node) (*Node, *Node) |
		func(*Graph) (*Node, *Node, *Node) |
		func([]*Node) (*Node, *Node, *Node) |
		func(*Node) (*Node, *Node, *Node) |
		func(*Node, *Node) (*Node, *Node, *Node) |
		func(*Node, *Node, *Node) (*Node, *Node, *Node) |
		func(*Node, *Node, *Node, *Node) (*Node, *Node, *Node) |
		func(*Node, *Node, *Node, *Node, *Node) (*Node, *Node, *Node) |
		func(*Node, *Node, *Node, *Node, *Node, *Node) (*Node, *Node, *Node) |
		func(*Graph) []*Node |
		func([]*Node) []*Node |
		func(*Node) []*Node |
		func(*Node, *Node) []*Node |
		func(*Node, *Node, *Node) []*Node |
		func(*Node, *Node, *Node, *Node) []*Node |
		func(*Node, *Node, *Node, *Node, *Node) []*Node |
		func(*Node, *Node, *Node, *Node, *Node, *Node) []*Node
}

// ExecGraphFnOneOutput are ExecGraphFn functions that return only one result.
// See MustExecOnce.
//
//nolint:goimports,golines // See https://youtrack.jetbrains.com/issue/GO-19556/Goland-formatter-gofmt
type ExecGraphFnOneOutput interface {
	func(*Graph) *Node |
		func([]*Node) *Node |
		func(*Node) *Node |
		func(*Node, *Node) *Node |
		func(*Node, *Node, *Node) *Node |
		func(*Node, *Node, *Node, *Node) *Node |
		func(*Node, *Node, *Node, *Node, *Node) *Node |
		func(*Node, *Node, *Node, *Node, *Node, *Node) *Node
}

// SideParamsFn is a function that sets side parameters during execution
// for Graphs that defines those. Typically, this is used to set the variables of a model.
type SideParamsFn func(graph *Graph, inputBuffers []backends.Buffer, donate []bool) error

// LoggerFn is the function used to log nodes marked for logging.
// It is called after the Exec method, with the list of messages and corresponding values of the evaluated nodes.
type LoggerFn func(graph *Graph, messages []string, values []*tensors.Tensor, nodes []NodeId)

// Exec creates and executes computation graphs as needed based on the inputs shapes.
//
// It simplifies the process of executing a graph building
// function with real values. For example, assume you wrote:
//
//	func L2Norm(x *Node) *Node {
//		return Sqrt(ReduceAllSum(Mul(x, x)))
//	}
//
// To use it with actual values (tensors.Tensor's), one needs to build
// the computation graph for the specific shape of x, and then execute it.
//
// While this is straight forward, it's lots of boilerplate code -- JIT compilation makes things
// faster, but it imposes some bureaucracy.
//
// With Exec one can do:
//
//	var l2NormExec = MustNewExec(backends.New(), L2Norm)
//	x0 := []float32{2}
//	fmt.Printf("L2Norm(%v) = %s\n", x0, l2NormExec.MustExec1(x0)) // -> 2
//	x1 := []float64{4, 3}
//	fmt.Printf("L2Norm(%v) = %s\n", x1, l2NormExec.MustExec1(x1))  // -> 5
//
// Notice that both calls to l2NormExec.MustExec1 will need to create different
// graphs (because they have different input shapes).
// These JIT-compiled computation graphs are cached,
// and if the same shapes are used in MustExec1 again, the cached version is reused.
//
// For "ergonomics", we provide variations of the same API, depending on if you want errors to be returned,
// or simply panic, and on how many outputs you expect:
//
//   - Use NewExec or MustNewExec to create new executors.
//   - Exec returns a slice of outputs and errors. Exec1 to Exec4 return 1 to 4 outputs exactly, plus an error.
//   - MustExec returns a slice of outputs and panics on error. MustExec1 to MustExec4 1 to 4 outputs exactly, and panics on error.
//   - ExecOnce and MustExecOnce merge NewExec (or MustNewExec), Exec1 (or MustExec1) and then Finalize methods in one convenient call,
//     for when executing the graph only once.
//     There are also ExecOnceN and MustExecOnceN variations that return a slice of outputs.
//
// If there are no inputs for the graph function (for instance, for some initialization function), then
// one needs to take a *Graph as the first parameter of the graph function (graphFn).
//
// Inside graph building functions, one is expected to panic with an error.
// The executor captures (recovers) from panics and returns the error.
//
// Example using `ExecOnce`, which compiles the graph and immediately runs it.
//
//	iotaMatrix, err := ExecOnce(backend, func (g *Graph) *Node {
//		return IotaFull(g, shapes.Make(dtype.Float32, 3, 3))
//	})
//	if err != nil { ... }
//	fmt.Printf("IotaFull(3x3 matrix, float32)=%v\n", iotaMatrix) // -> [[0 1 2] [3 4 5] [6 7 8]]
//
// Example using `ExecOnceN`, which returning 2 values:
//
//	iotaMatrices, err := ExecOnceN(backend, func (g *Graph) []*Node {
//		return []*Node{
//			IotaFull(g, shapes.Make(dtype.Float32, 2, 2)),
//			IotaFull(g, shapes.Make(dtype.Float32, 3, 3))}
//	})
//
// # Considerations of inputs of different shapes:
//
// The need to build different graphs for different shapes can be expensive when the inputs' shapes vary a lot.
// The usual solution is to use shapes with dimensions in a power scale (for instance, powers of 2) and
// use padding and masking of tensors for unused slices of the input.
//
// For safety concerns, there are a maximum number of different instantiations of the graph.
// It can be set or disabled with SetMaxCache.
//
// # Logging:
//
// The executor can log nodes marked for logging during execution, which is very handy for debugging, plotting, etc.
// One marks a computation (any Node) for logging with Node.SetLogged (or Node.SetLoggedf).
// And one can have specialized loggers by setting use Exec.SetNodeLogger: they can be set in cascading fashion,
// so there can be different logger types (printing, plotting, NaN detection, etc.) being used simultaneously.
//
// # Distributed execution:
//
//   - User controlled (the default): the user can choose the device to use for execution with WithDevice.
//     Only one device is used (see Exec.WithDevice), but the user can create multiple executors for different
//     devices and execute code in parallel. No communication (distributed operations) available among the executors.
//   - SPMD (Single Program Multiple Data): automatically replicates the computation across devices, see Exec.SPMD.
//     The devices are organized in an arbitrary mesh. Usually, it's a 1D list of devices, but they can be organized
//     in multiple axes. One uses Graph.Distributed() for the synchronization operations
//     (like AllReduce, AllGather, etc.), which need to be organized by the user.
type Exec struct {
	name    string
	backend backends.Backend

	distStrategy     distributed.Strategy
	meshes           []*distributed.DeviceMesh // For distributed.SPMD or distributed.AutoSharding.
	deviceAssignment []backends.DeviceNum
	numDevices       int

	// graphFn is fixed for all entries in the cache.
	graphFn any

	// numInputs, numOutputs of graphFn, not counting extra logged nodes, which may vary
	// per instance of the graph (entry in the cache).
	// This is relative if inputAsSlice is set to true, in which case graphFn's number of inputs/outputs vary
	// per cache entry.
	numInputs, numOutputs int

	inputAsSlice, outputAsSlice bool
	inputIsGraph                bool
	inputShardingSpecs          []*distributed.ShardingSpec
	outputShardingSpecs         []*distributed.ShardingSpec

	// MaxCacheSize: if more than these different graph instantiations are
	// created, Exec starts returning errors in MustExec.
	maxCacheSize int

	// setSideParams for graphs that take them.
	setSideParams SideParamsFn
	loggerFn      LoggerFn

	// Protects cache structure.
	cacheMu sync.Mutex
	cache   []*execGraphCacheEntry
}

// execGraphCacheEntry: no hashing, just a simple list. This is faster
// for smaller tables. TODO: add a hashtable for cases with large caches.
type execGraphCacheEntry struct {
	argsShapes        []shapes.Shape
	graph             *Graph
	numGraphFnOutputs int      // Number
	numAllOutputs     int      // Number of flattened outputs for this graph, including logged nodes.
	loggedMessages    []string // Messages for logged nodes.
	loggedNodeIDs     []NodeId
}

// DefaultExecMaxCacheSize is the value used to initialize the max cache size of new Exec objects.
const DefaultExecMaxCacheSize = 32

// NewExecAny constructs an Exec object that uses the given graphFn to build
// computation graphs.
//
// The `graphFn` can take only *Node parameters as input and returns one or more *Node.
// Except if there are no inputs, in which case graphFn needs to take a *Graph as the first parameter.
//
// Please use NewExec if possible, it adds a compile-time check for most valid signatures of graphFn.
//
// It returns an error if the inputs are invalid.
func NewExecAny(backend backends.Backend, graphFn any) (*Exec, error) {
	funcName := runtime.FuncForPC(reflect.ValueOf(graphFn).Pointer()).Name()
	e := &Exec{
		backend:      backend,
		name:         fmt.Sprintf("Exec:%s", funcName),
		graphFn:      graphFn,
		maxCacheSize: DefaultExecMaxCacheSize,
		loggerFn:     DefaultNodeLogger,
		numDevices:   1,
	}
	if err := e.parseGraphFn(); err != nil {
		return nil, err
	}
	return e, nil
}

// NewExec constructs an Exec object that uses the given graphFn to build computation graphs.
//
// graphFn should take *Node as an input and return a *Node -- except if there are no (Node) inputs,
// in which case it should take a single *Graph input.
//
// It's a wrapper to the NewExecAny, but it uses generics to add a compile-time check for valid graphFn signature.
func NewExec[F ExecGraphFn](backend backends.Backend, graphFn F) (*Exec, error) {
	return NewExecAny(backend, graphFn)
}

// parseGraphFn figures out and validates the inputs and outputs of the graphFn for the Exec object.
// This should be called in the constructor.
//
//nolint:gocognit
func (e *Exec) parseGraphFn() error {
	graphFnT := reflect.TypeOf(e.graphFn)
	e.numInputs = graphFnT.NumIn()
	e.numOutputs = graphFnT.NumOut()

	// Verify parameters.
	if graphFnT.Kind() != reflect.Func {
		return errors.Errorf("graphFn must be a function")
	}

	var node *Node
	nodeType := reflect.TypeOf(node)
	var tmpGraph *Graph
	graphType := reflect.TypeOf(tmpGraph)

	if graphFnT.NumIn() < 1 || graphFnT.NumOut() < 1 {
		// It requires at least one input and one output.
		return errors.Errorf("not enough input (%d)/output (%d) parameters, both need to be > 0",
			graphFnT.NumIn(), graphFnT.NumOut())
	}
	for ii := range graphFnT.NumIn() {
		if graphFnT.In(ii).Kind() == reflect.Slice && graphFnT.In(ii).Elem() == nodeType {
			if graphFnT.NumIn() != 1 {
				return errors.Errorf("[]*Node parameters are only accepted as input if they are the only input,"+
					" got function type %s instead", graphFnT)
			}
			e.inputAsSlice = true
			break
		}
		if graphFnT.In(ii) == graphType {
			if graphFnT.NumIn() != 1 {
				return errors.Errorf("*Graph parameter only accepted as input if they are the only input, got "+
					"function type %s instead", graphFnT)
			}
			e.inputIsGraph = true
			e.numInputs = 0
			break
		}
		if graphFnT.In(ii) != nodeType {
			return errors.Errorf("input parameter %d is not of type *Node or []*Node", ii)
		}
	}
	for ii := range graphFnT.NumOut() {
		if graphFnT.Out(ii).Kind() == reflect.Slice && graphFnT.Out(ii).Elem() == nodeType {
			if graphFnT.NumOut() != 1 {
				return errors.Errorf("[]*Node parameters are only accepted as output if they are the only output,"+
					" got function type %s instead", graphFnT)
			}
			e.outputAsSlice = true
			break
		}
		if graphFnT.Out(ii) != nodeType {
			return errors.Errorf("output parameter %d is not of type *Node", ii)
		}
	}
	return nil
}

// WithName sets the name of Exec, used to provide the name to graphs created.
// This should be called before any invocations of MustExec().
// It returns a reference to itself so calls can be cascaded.
func (e *Exec) WithName(name string) *Exec {
	e.name = name
	return e
}

// WithDeviceAssignment specifies which concrete devices to use when compiling computation graphs.
//
// These must be valid numbers for the backend and must match the number of devices of the
// largest mesh given to WithAutoSharding or WithSPMD, or one fixed device for non-portable single-device
// execution.
//
// The default assignment is simply using the devices in the order they were added to the backend
// (sequential DeviceNum values, starting from 0).
//
// For single-device execution (distributed strategy "None"), the default is to make it portable.
// If the backend supports that, it can be executed in any device with ExecOnDevice().
func (e *Exec) WithDeviceAssignment(devices []backends.DeviceNum) *Exec {
	e.deviceAssignment = devices
	return e
}

// DeviceAssignment returns the current device assignment used by this Exec.
// It returns nil if none was provided.
func (e *Exec) DeviceAssignment() []backends.DeviceNum {
	return e.deviceAssignment
}

// DistributionStrategy returns the distribution strategy used by this Exec.
//
// The default is distributed.None, which means that graphs constructed by this Exec will not be distributed,
// and it will be executed on the device specified by Exec.WithDevice (defaults to 0).
func (e *Exec) DistributionStrategy() distributed.Strategy {
	return e.distStrategy
}

// SPMD sets the distribution strategy to SPMD, which means that graphs constructed by this Exec will be replicated
// across the devices specified in the mesh.
//
// A nil mesh will cause a panic.
//
// It returns a reference to itself, so configuration calls can be cascaded.
func (e *Exec) SPMD(mesh *distributed.DeviceMesh) *Exec {
	e.distStrategy = distributed.SPMD
	e.meshes = []*distributed.DeviceMesh{mesh}
	e.numDevices = mesh.NumDevices()
	return e
}

// AutoSharding sets the distribution strategy to AutoSharding and records the meshes that will be used in the
// computation graph(s).
//
// A nil mesh will cause a panic.
//
// It returns a reference to itself, so configuration calls can be cascaded.
func (e *Exec) AutoSharding(meshes ...*distributed.DeviceMesh) *Exec {
	e.distStrategy = distributed.AutoSharding
	e.meshes = slices.Clone(meshes)
	e.numDevices = 0
	for _, mesh := range meshes {
		e.numDevices = max(e.numDevices, mesh.NumDevices())
	}
	return e
}

// Meshes returns the slice of currently configured meshes.
// It returns nil if no meshes were provided (e.g., for non-distributed execution).
func (e *Exec) Meshes() []*distributed.DeviceMesh {
	return e.meshes
}

// WithInputShardingSpecs sets the sharding specs for the inputs.
//
// This is used for distributed computations with AutoSharding.
//
// If the function takes variable inputs (`[]*Node`), then the last spec provided is used for all remaining inputs.
//
// It returns a reference to itself so calls can be cascaded.
func (e *Exec) WithInputShardingSpecs(specs ...*distributed.ShardingSpec) *Exec {
	e.inputShardingSpecs = specs
	return e
}

// WithOutputShardingSpecs sets the sharding specs for the outputs.
//
// If the function takes variable inputs (`[]*Node`), then the last spec provided is used for all remaining inputs.
//
// This is used for distributed computations with AutoSharding.
//
// It returns a reference to itself so calls can be cascaded.
func (e *Exec) WithOutputShardingSpecs(specs ...*distributed.ShardingSpec) *Exec {
	e.outputShardingSpecs = specs
	return e
}

// NumDevices returns the number of devices used by the current strategy.
func (e *Exec) NumDevices() int {
	return e.numDevices
}

// Name returns the Exec name, a string used as a prefix for Graph construction.
func (e *Exec) Name() string {
	return e.name
}

// SetMaxCache sets the maximum size of the cache.
// Set it to -1 to have unlimited cache size.
// It returns a reference to itself so calls can be cascaded.
func (e *Exec) SetMaxCache(maxCacheSize int) *Exec {
	e.cacheMu.Lock()
	defer e.cacheMu.Unlock()
	e.maxCacheSize = maxCacheSize
	return e
}

// SetSideParamsHook configures a function to be called just before executing a graph, so it can set extra parameters.
//
// Mostly, this is for internal use and end-users will not likely need this. The context.Exec object uses this to pass
// the variable values as side inputs to the graph.
//
// Exec takes care of creating parameters (with graph.Parameter) for every value passed to Exec.Exec before
// calling the graph building function (the graph building function is executed only the first time, after the
// graph is compiled, it is re-used for future executions).
//
// But a graph building function may want to create extra parameters itself (with graph.Parameter), which we call
// "side parameters".
//
// The values to feed these "side parameters" are not passed to Exec.MustExec, but instead set with a SideParamsFn, which
// is configured here.
//
// SideParamsFn is called after the graph is already built, just before the execution.
// It is passed with a slice of the backend.Buffer to be fed to the graph execution.
// The side parameters in this slice will be left nil, and it's expected that SideParamsFn will set
// them to the appropriate input.
//
// Notice: len(inputBuffers) = len(donate) = g.NumDevices() * g.NumParameters(), organized by device first.
// The SideParamsFn needs to set the last parameters (currently nil) for each device,
// in the order they were added to the backend.
//
// It also includes the boolean map of the inputs to donate, which SideParamsFn
// can set accordingly (for the side parameters).
func (e *Exec) SetSideParamsHook(fn SideParamsFn) *Exec {
	e.setSideParams = fn
	return e
}

// SetNodeLogger with the function to be called for the nodes marked for logging during execution.
// If set to nil, nothing will be logged.
func (e *Exec) SetNodeLogger(loggerFn LoggerFn) {
	e.loggerFn = loggerFn
}

// GetNodeLogger returns the currently registered LoggerFn.
func (e *Exec) GetNodeLogger() LoggerFn {
	return e.loggerFn
}

// ExecWithGraphOnDevice is a version of ExecGraph that runs on the given device by default.
//
// deafultDevice is used for single-device computations that are portable (no fixed device assignment set
// WithDeviceAssignment). Otherwise, it is ignored.
func (e *Exec) ExecWithGraphOnDevice(defaultDevice backends.DeviceNum, args ...any) (
	[]*tensors.Tensor, *Graph, error) {
	outputs, g, err := e.compileAndExecute(true, defaultDevice, args...)
	if err != nil {
		return nil, nil, err
	}
	return outputs, g, nil
}

// PreCompile will build the computation graph and compile it but not yet execute.
// Useful when one wants to measure the time separately, from graph compilation and its execution.
//
// Notice, this will include the time to convert args to tensors. If you want to isolate that time,
// pre-convert args to tensors first.
func (e *Exec) PreCompile(args ...any) error {
	_, _, err := e.compileAndExecute(false, 0, args...)
	return err
}

// unwrapListOfTensors converts something like []any{[]*tensors.Tensor{t1, t2, ...}} to []any{t1, t2,...}.
// If args is something else, it remains untouched.
func unwrapListOfTensors(args []any) []any {
	if len(args) != 1 {
		return args
	}
	if list, ok := args[0].([]*tensors.Tensor); ok {
		return xslices.Map(list, func(x *tensors.Tensor) any { return x })
	}
	// Otherwise, process as usual.
	return args
}

// unwrapDistributedTensor converts something like []any{distributed.Tensor0, distributed.Tensor1, ...} to
// []any{shard0_0, shard0_1, ... , shard1_0, shard1_1, ... , shardN_0, shardN_1, ...},
// where shardI_J is the shard J of the distributed.Tensor I.
func unwrapDistributedTensors(numDevices int, args []any) ([]any, error) {
	if len(args) == 0 {
		return args, nil
	}
	_, ok := args[0].(*distributed.Tensor)
	if !ok {
		// No distributed.Tensors, so no unwrapping needed.
		return args, nil
	}
	unwrapped := make([]any, numDevices*len(args))
	for argIdx, arg := range args {
		dTensor, ok := arg.(*distributed.Tensor)
		if !ok {
			return nil, errors.Errorf(
				"argument #%d is not a distributed.Tensor -- if using distributed.Tensor as input, all arguments "+
					"must be distributed.Tensors, or none of them", argIdx)
		}
		shards := dTensor.Shards()
		if len(shards) != numDevices {
			return nil, errors.Errorf(
				"distributed.Tensor #%d has %d shards, but graph has %d devices and those many shards  are expected",
				argIdx, len(shards), numDevices)
		}
		for shardIdx, shard := range shards {
			unwrapped[shardIdx*len(args)+argIdx] = shard
		}
	}
	return unwrapped, nil
}

// validateArgs to number of expected number of inputs and devices.
//
// deafultDevice is used for single-device computations that are portable (no fixed device assignment).
// Otherwise, it is ignored.
func (e *Exec) validateArgs(args []any) error {
	if len(args) == 0 {
		if !e.inputAsSlice && e.numInputs != 0 {
			return errors.Errorf("no arguments passed to Exec for %q, but graphFn expects %d inputs",
				e.Name(), e.numInputs)
		}
		// Having no arguments is valid for the graph.
		return nil
	}

	// Check that args are a multiple of the number of devices involved.
	numDevices := e.numDevices
	if numDevices > 1 && len(args)%numDevices != 0 {
		return errors.Errorf("number of arguments for execution (%d) doesn't match number of devices (%d) for %q",
			len(args), numDevices, e.Name())
	}

	// Verify that numArgsPerDevice matches the number of inputs expected by the graphFn.
	numArgsPerDevice := len(args) / numDevices
	if !e.inputAsSlice && numArgsPerDevice != e.numInputs {
		return errors.Errorf(
			"# of arguments to call (#args=%d, %d per device) don't match # arguments to the graph function "+
				"(#args=%d) for %q",
			len(args), numArgsPerDevice, e.numInputs, e.Name())
	}
	return nil
}

// convertArgsToBuffers converts the arguments as any type to buffers, taking into account the device assignment.
func (e *Exec) convertArgsToBuffers(args []any, defaultDevice backends.DeviceNum) (
	argsAsBuffers []backends.Buffer, argsShapes []shapes.Shape, argsDonate []bool, err error) {
	// Convert args to buffers: care is taken so we move each value to the correct device.
	// Note there may be more parameters, set with Exec.setSideParams later.
	argsAsBuffers = make([]backends.Buffer, len(args))
	argsShapes = make([]shapes.Shape, len(args))
	argsDonate = make([]bool, len(args))
	numDevices := e.numDevices
	argsPerDevice := len(args) / numDevices
	var argIdx int
	argDeviceNum := defaultDevice
	for argDeviceIdx := range numDevices {
		if numDevices > 1 {
			if len(e.deviceAssignment) <= argDeviceIdx {
				// If deviceAssignment is not given, we assume an f(idx) = idx assignment of devices.
				argDeviceNum = backends.DeviceNum(argDeviceIdx)
			} else {
				argDeviceNum = e.deviceAssignment[argDeviceIdx]
			}
		}
		for range argsPerDevice {
			// TODO: set argDeviceNum according to device assignment.
			arg := args[argIdx]
			err := exceptions.TryCatch[error](func() {
				argsAsBuffers[argIdx], argsShapes[argIdx], argsDonate[argIdx] = anyToDeviceBuffer(
					e.backend,
					argDeviceNum,
					arg,
				)
			})
			if err != nil {
				return nil, nil, nil, errors.WithMessagef(
					err, "Failed to convert argument #%d of %d to device(%d) -- type %T: %v",
					argIdx, len(args), argDeviceNum, args[argIdx], args[argIdx])
			}
			argIdx++
		}
	}
	return argsAsBuffers, argsShapes, argsDonate, nil
}

// expandArgsToTotalParams expands the arguments to the total number of parameters for the graph.
// If the graph has more parameters than the arguments, it will add new parameters to the arguments.
//
// This happens if the graph building function created extra graph.Parameter nodes, that Exec is not aware of.
// The package ml/context does that for variables used in during model building.
//
// The slices are organized in "deviceNum major" order:
// [device0_params..., device1_params..., device2_params..., ...]
// So after growing, we need to shuffle each device's parameters to the start of their expanded row.
func (e *Exec) expandArgsToTotalParams(g *Graph, argsAsBuffers []backends.Buffer, argsDonate []bool) (
	[]backends.Buffer, []bool) {
	newParamsPerDevice := g.NumParameters()
	totalParams := newParamsPerDevice * e.numDevices
	if totalParams <= len(argsAsBuffers) {
		return argsAsBuffers, argsDonate
	}

	oldParamsPerDevice := len(argsAsBuffers) / e.numDevices
	numNew := totalParams - len(argsAsBuffers)
	argsAsBuffers = slices.Grow(argsAsBuffers, numNew)
	argsAsBuffers = argsAsBuffers[:totalParams]
	argsDonate = slices.Grow(argsDonate, numNew)
	argsDonate = argsDonate[:totalParams]

	// Shuffle each device's parameters to the start of their row in the expanded slice.
	// Work backwards to avoid overwriting data.
	for deviceIdx := e.numDevices - 1; deviceIdx >= 0; deviceIdx-- {
		oldStart := deviceIdx * oldParamsPerDevice
		newStart := deviceIdx * newParamsPerDevice

		// Copy existing parameters to their new position.
		// Note: Go's copy handles overlapping slices correctly (uses memmove semantics).
		copy(argsAsBuffers[newStart:newStart+oldParamsPerDevice], argsAsBuffers[oldStart:oldStart+oldParamsPerDevice])
		copy(argsDonate[newStart:newStart+oldParamsPerDevice], argsDonate[oldStart:oldStart+oldParamsPerDevice])

		// Clear the new slots (and old positions that are now in the "new slots" region).
		for i := newStart + oldParamsPerDevice; i < newStart+newParamsPerDevice; i++ {
			argsAsBuffers[i] = nil
			argsDonate[i] = false
		}
	}
	return argsAsBuffers, argsDonate
}

// compileAndExecute compiles a graph for arguments and optionally executes it.
//
// deafultDevice is used for single-device computations that are portable (no fixed device assignment set
// WithDeviceAssignment). Otherwise, it is ignored.
func (e *Exec) compileAndExecute(execute bool, defaultDevice backends.DeviceNum, args ...any) (
	[]*tensors.Tensor, *Graph, error) {
	args = unwrapListOfTensors(args)
	var err error
	args, err = unwrapDistributedTensors(e.numDevices, args)
	if err != nil {
		return nil, nil, err
	}
	err = e.validateArgs(args)
	if err != nil {
		return nil, nil, err
	}
	numDevices := e.numDevices
	argsAsBuffers, argsShapes, argsDonate, err := e.convertArgsToBuffers(args, defaultDevice)
	if err != nil {
		return nil, nil, err
	}

	// Get or build the graph.
	var entry *execGraphCacheEntry
	err = exceptions.TryCatch[error](func() {
		entry = e.findOrCreateGraph(argsShapes)
	})
	if err != nil {
		return nil, nil, err
	}
	if entry == nil {
		return nil, nil, errors.Errorf(
			"maximum cache size of %d reached for %q, cannot create another graph -- "+
				"a new computation graph needs to be created+compiled for each different shape of "+
				"the input, consider using padding, or if this is not a concern change "+
				"the cache size with executable.SetMaxCache()", e.maxCacheSize, e.Name())
	}
	g := entry.graph

	// Now that the graph is created, we know the exact number of parameters: if the graph building function created
	// new graph.Parameter, we may need to include those in our argsAsBuffer and argsDonate accordingly.
	argsAsBuffers, argsDonate = e.expandArgsToTotalParams(g, argsAsBuffers, argsDonate)

	// The new parameters (if any) created are still nil and need to be set. This is done by a "SideParamsFn",
	// configured by Exec.SetSideParamsHooks.
	if e.setSideParams != nil {
		err := e.setSideParams(g, argsAsBuffers, argsDonate)
		if err != nil {
			return nil, nil, errors.WithMessagef(err, "failed to set side parameters for graph %q", g.Name())
		}
	}

	// Check all parameters were set.
	numParams := g.NumParameters()
	for ii, t := range argsAsBuffers {
		if t == nil {
			paramIdx := ii % numParams
			deviceIdx := ii / g.NumParameters()
			if numDevices == 1 {
				return nil, nil, errors.Errorf("parameter #%d (%q) is nil or invalid, maybe a variable value not set as a "+
					"parameter, cannot execute the graph %q",
					paramIdx, g.GetParameterByHandle(ParameterHandle(paramIdx)).GetParameterName(), g.Name())
			}
			return nil, nil, errors.Errorf("parameter #%d (%q) for device #%d is nil or invalid, maybe a variable value not set as a "+
				"parameter, cannot execute the graph %q",
				paramIdx, g.GetParameterByHandle(ParameterHandle(paramIdx)).GetParameterName(), deviceIdx, g.Name())
		}
	}

	// To only pre-compile the graph, return early.
	if !execute {
		return nil, g, nil
	}

	// Execute graph: outputs will have (numDevice * entry.numAllOutputs) outputs.
	var outputs []*tensors.Tensor
	err = exceptions.TryCatch[error](func() {
		outputs = g.RunWithBuffers(argsAsBuffers, argsDonate, defaultDevice)
	})
	if err != nil {
		return nil, nil, errors.WithMessagef(err, "failed to execute graph %q", g.Name())
	}
	if len(outputs) != entry.numAllOutputs*e.numDevices {
		return nil, nil, errors.Errorf("expected %d * %d (numDevices) = %d outputs from graph %q, got %d",
			e.numOutputs, e.numDevices, e.numOutputs*e.numDevices, g.Name(), len(outputs))
	}

	// Call the logger on logged nodes, even if no node is marked for logging (it serves as a hook).
	numGraphFnOutputs := entry.numAllOutputs - len(entry.loggedMessages)
	if e.loggerFn != nil {
		var loggerOutputs []*tensors.Tensor
		if len(entry.loggedMessages) > 0 {
			// If the computation is distributed, the Logged outputs are all replicated,
			// so we only use (slice) the ones returned in the first device.
			loggerOutputs = outputs[numGraphFnOutputs:entry.numAllOutputs]
		}
		// The logger is also called for zero logged messages.
		e.loggerFn(g, entry.loggedMessages, loggerOutputs, entry.loggedNodeIDs)
	}

	// Remove logged messages from outputs, we need to take slices for each device:
	if len(entry.loggedMessages) == 0 {
		// Easiest case: no logged messages, no slice needed.
		return outputs, g, nil
	}
	if numDevices == 1 {
		// No need to rebuild a new array:
		return outputs[:numGraphFnOutputs], g, nil
	}
	graphFnOutputs := make([]*tensors.Tensor, numDevices*numGraphFnOutputs)
	for deviceIdx := range numDevices {
		copy(graphFnOutputs[deviceIdx*numGraphFnOutputs:(deviceIdx+1)*numGraphFnOutputs],
			outputs[deviceIdx*entry.numAllOutputs:deviceIdx*entry.numAllOutputs+numGraphFnOutputs])
	}
	return graphFnOutputs, g, nil
}

// createAndCacheGraph creates and compiles the graph for the arguments with the given
// shapes. It creates and stores a cache entry for it and returns it.
// Returns nil if the cache size is >= MaxCacheSize.
// Should be called with cacheMu locked.
func (e *Exec) createAndCacheGraph(argsShapes []shapes.Shape) *execGraphCacheEntry {
	if e.maxCacheSize > 0 && len(e.cache) >= e.maxCacheSize {
		return nil
	}
	entry := &execGraphCacheEntry{graph: NewGraph(e.backend, fmt.Sprintf("%s#%d", e.name, len(e.cache)))}
	g := entry.graph
	switch e.distStrategy {
	case distributed.AutoSharding:
		err := g.SetAutoSharding(e.meshes...)
		if err != nil {
			panic(errors.WithMessagef(err, "failed to assign AutoSharding for graph %q", g.Name()))
		}
	case distributed.SPMD:
		err := g.SetSPMD(e.meshes[0])
		if err != nil {
			panic(errors.WithMessagef(err, "failed to assign SPMD for graph %q", g.Name()))
		}
	case distributed.None:
		// Nothing to do.
	}
	if e.deviceAssignment != nil {
		err := g.SetDeviceAssignment(e.deviceAssignment)
		if err != nil {
			panic(
				errors.WithMessagef(
					err,
					"failed to assign device assignment %v for graph %q",
					e.deviceAssignment,
					g.Name(),
				),
			)
		}
	}

	var argsV []reflect.Value
	var args []*Node

	// Calculate the actual args to use for graph creation:
	// If distributed with AutoSharding or SPMD, the argsShapes contains the shapes for ALL devices (concatenated).
	// We only want the shapes for the first device (or first replica).
	//
	// Note: AutoSharding uses a "logical" graph, so we need to adjust the shapes according to the sharding specs.
	// But first we filter out the shapes for other devices.
	argsShapesToUse := argsShapes
	if e.numDevices > 1 && (e.distStrategy == distributed.AutoSharding || e.distStrategy == distributed.SPMD) {
		// argsShapes is organized as [device0_args..., device1_args..., ...].
		// We only need the args for one device to build the graph.
		numArgsPerDevice := len(argsShapes) / e.numDevices
		argsShapesToUse = argsShapes[:numArgsPerDevice]
	}

	switch {
	case e.inputAsSlice:
		args = make([]*Node, 0, len(argsShapesToUse))
	case e.inputIsGraph:
		// Notice in this case len(argsShapesToUse) == 0
		argsV = []reflect.Value{reflect.ValueOf(g)}
	default:
		argsV = make([]reflect.Value, 0, len(argsShapesToUse))
	}
	for ii, shape := range argsShapesToUse {
		var spec *distributed.ShardingSpec
		if ii < len(e.inputShardingSpecs) {
			spec = e.inputShardingSpecs[ii]
		} else if e.inputAsSlice && len(e.inputShardingSpecs) > 0 {
			// Reuse last spec for variable length inputs.
			spec = e.inputShardingSpecs[len(e.inputShardingSpecs)-1]
		}
		if spec != nil && e.distStrategy == distributed.AutoSharding {
			// Adjust shape to logical shape.
			shape = spec.LogicalShapeForShard(shape)
		}

		arg := ShardedParameter(g, fmt.Sprintf("arg%d", ii), shape, spec)
		if e.inputAsSlice {
			args = append(args, arg)
		} else {
			argsV = append(argsV, reflect.ValueOf(arg))
		}
	}
	graphFnV := reflect.ValueOf(e.graphFn)
	if e.inputAsSlice {
		// If input is a slice of *Node, take argsV to be one parameter, the value of the slice.
		argsV = []reflect.Value{reflect.ValueOf(args)}
	}

	// Enumerate outputs from wrapped graphFn.
	start := time.Now()
	outputsV := graphFnV.Call(argsV)
	if klog.V(1).Enabled() {
		elapsed := time.Since(start)
		klog.Infof("Build graph time for %q: %s", e.name, elapsed)
	}

	var outputs []*Node
	if e.outputAsSlice {
		if len(outputsV) != 1 {
			exceptions.Panicf("graphFn for %q returned %d results, as opposed to simply a slice of nodes -- []*Node",
				e.Name(), len(outputsV))
		}
		var ok bool
		outputs, ok = outputsV[0].Interface().([]*Node)
		if !ok {
			exceptions.Panicf("graphFn for %q did not return a []*Node, instead it returned %T!?",
				e.Name(), outputsV[0].Interface())
		}
		entry.numGraphFnOutputs = len(outputs)
	} else {
		entry.numGraphFnOutputs = e.numOutputs // Fixed number of outputs.
		if len(outputsV) != e.numOutputs {
			exceptions.Panicf(
				"graphFn for %q returned %d results, as opposed to the actual returned %d results -- []*Node",
				e.Name(), len(outputsV), e.numOutputs)
		}
		outputs = make([]*Node, 0, len(outputsV))
		for i, outV := range outputsV {
			outputNode, ok := outV.Interface().(*Node)
			if !ok {
				exceptions.Panicf("graphFn for %q did not return a *Node for output #%d, instead it returned %T!?",
					e.Name(), i, outV.Interface())
			}
			outputs = append(outputs, outputNode)
		}
	}

	// Append logged nodes as outputs.
	loggedNodes := g.LoggedNodes()
	entry.loggedMessages = make([]string, 0, len(loggedNodes))
	entry.loggedNodeIDs = make([]NodeId, 0, len(loggedNodes))
	for _, node := range loggedNodes {
		outputs = append(outputs, node)
		entry.loggedMessages = append(entry.loggedMessages, node.LogMessage())
		entry.loggedNodeIDs = append(entry.loggedNodeIDs, node.Id())
	}

	// Append ShardingSpec for logged nodes: they are always replicated.
	outputShardingSpecs := e.outputShardingSpecs
	if len(outputShardingSpecs) > 0 {
		if len(outputShardingSpecs) < entry.numGraphFnOutputs {
			// For variable length outputs, we repeat the last spec.
			lastSpec := outputShardingSpecs[len(outputShardingSpecs)-1]
			for len(outputShardingSpecs) < entry.numGraphFnOutputs {
				outputShardingSpecs = append(outputShardingSpecs, lastSpec)
			}
		}
		// Logged nodes spec is "replicated" (nil).
		for range len(loggedNodes) {
			outputShardingSpecs = append(outputShardingSpecs, nil)
		}
	}

	// Compile graph.
	g.CompileWithSharding(outputs, outputShardingSpecs)
	entry.argsShapes = make([]shapes.Shape, len(argsShapes))
	copy(entry.argsShapes, argsShapes)
	entry.numAllOutputs = len(outputs)
	e.cache = append(e.cache, entry)
	return entry
}

// findOrCreateGraph returns the graph for the given arguments shapes: either from cache or by creating a new one.
// if no cache entry exists.
func (e *Exec) findOrCreateGraph(argsShapes []shapes.Shape) *execGraphCacheEntry {
	e.cacheMu.Lock()
	defer e.cacheMu.Unlock()

LoopCache:
	for _, entry := range e.cache {
		if len(argsShapes) != len(entry.argsShapes) {
			continue
		}
		for ii, shape := range argsShapes {
			if !shape.Equal(entry.argsShapes[ii]) {
				continue LoopCache
			}
		}
		return entry
	}

	// No graph in cache, create a new one.
	return e.createAndCacheGraph(argsShapes)
}

// Finalize clears the cache, finalizing the compiled graphs. The Exec object shouldn't be
// used after that.
func (e *Exec) Finalize() {
	e.cacheMu.Lock()
	defer e.cacheMu.Unlock()

	for _, entry := range e.cache {
		entry.graph.Finalize()
		entry.graph = nil
	}
	e.cache = e.cache[:0]
}

// DefaultNodeLogger for nodes marked to be logged. It prints the message and
// the node value for each logged node.
//
// It accepts special prefixes on message name that affects the printing:
//
//   - "#full ": prints full tensor value (as opposed to abbreviated).
//
//nolint:forbidigo  // This is an exported printing function.
func DefaultNodeLogger(g *Graph, messages []string, values []*tensors.Tensor, nodes []NodeId) {
	if len(messages) == 0 {
		return
	}
	fmt.Printf("DefaultNodeLogger(Graph %q):\n", g.Name())
	for ii, msg := range messages {
		if strings.HasPrefix(msg, "#full ") {
			fmt.Printf("\t(Node #%d) %s: %s\n", nodes[ii], msg[6:], values[ii].GoStr())
			continue
		}
		fmt.Printf("\t(Node #%d) %s: %s\n", nodes[ii], msg, values[ii])
	}
}
